/**
 * Copyright 2012 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.cluster.mom;

import org.apache.commons.math3.random.RandomGenerator;

import com.google.common.base.Preconditions;

import edu.byu.nlp.cluster.BasicProbabilisticClustering;
import edu.byu.nlp.cluster.Dataset;
import edu.byu.nlp.cluster.ProbabilisticClusterer;
import edu.byu.nlp.cluster.ProbabilisticClustering;
import edu.byu.nlp.cluster.ProbabilisticModel;
import edu.byu.nlp.cluster.em.ParameterInitializer;
import edu.byu.nlp.data.SparseFeatureVector;
import edu.byu.nlp.math.optimize.IterativeOptimizer;
import edu.byu.nlp.math.optimize.IterativeOptimizer.Optimizable;
import edu.byu.nlp.math.optimize.IterativeOptimizer.ReturnType;
import edu.byu.nlp.math.optimize.ValueAndObject;
import edu.byu.nlp.stats.DirichletDistribution;
import edu.byu.nlp.stats.RandomGenerators;
import edu.byu.nlp.util.Enumeration;
import edu.byu.nlp.util.Iterables2;

/**
 * @author rah67
 *
 */
public class CollapsedClusterer implements ProbabilisticClusterer {

	private final double alpha;
	private final double beta;
	private final IterativeOptimizer optimizer;
	private final ParameterInitializer<MoMParameters> initializer;
	private final RandomGenerator rnd;
	private final boolean isHard;

	/**
	 * Constructs a Clusterer based on the collapsed distribution. If isHard is false, then sampling is performed,
	 * otherwise, coordinate-wise gradient ascent is performed.
	 */
	public CollapsedClusterer(double alpha, double beta, IterativeOptimizer optimizer,
			ParameterInitializer<MoMParameters> initializer, boolean isHard, RandomGenerator rnd) {
		Preconditions.checkArgument(alpha > 0.0, "alpha (%s) must be strictly greater than 0.0", alpha);
		Preconditions.checkArgument(beta > 0.0, "beta (%s) must be strictly greater than 0.0", beta);
		Preconditions.checkNotNull(optimizer);
		Preconditions.checkNotNull(initializer);
		Preconditions.checkNotNull(rnd);

		this.alpha = isHard ? alpha - 1 : alpha;
		this.beta = isHard ? beta - 1 : beta;
		this.optimizer = optimizer;
		this.initializer = initializer;
		this.isHard = isHard;
		this.rnd = rnd;
	}

	private static final class CollapsedOptimizable implements Optimizable<CollapsedParameters> {

		private final Assigner assigner;
		
		public CollapsedOptimizable(Assigner assigner) {
			this.assigner = assigner;
		}
		
		/** {@inheritDoc} */
		@Override
		public ValueAndObject<CollapsedParameters> computeNext(CollapsedParameters prevParams) {
			prevParams.nextAssignments(assigner);
			return new ValueAndObject<CollapsedParameters>(prevParams.logJoint(), prevParams);
		}
		
	}

	public ProbabilisticClustering cluster(Dataset data, int numClusters) {
		int[] initialY = initializeAssignments(data, numClusters);
		CollapsedParameters sample = CollapsedParameters.newInstance(data, initialY, numClusters, alpha, beta);
		
		CollapsedOptimizable optimizable = null;
		if (isHard) {
			new CollapsedOptimizable(new ModeAssigner());
		} else {
			new CollapsedOptimizable(new SampleAssigner(rnd));
		}
		ValueAndObject<CollapsedParameters> sampleAndLogJoint = optimizer.optimize(optimizable, ReturnType.LAST, true,
				sample);
		
		// FIXME : what's the predictive distribution?
		return new BasicProbabilisticClustering(data, numClusters, sampleAndLogJoint.getObject().getY(),
				sampleModel(sampleAndLogJoint.getObject()), sampleAndLogJoint.getValue());
	}
	
	private int[] initializeAssignments(Dataset data, int numClusters) {
		MoMParameters params = initializer.initialize(data, numClusters);
		int[] initialY = new int[data.unlabeledData().size()];
		for (Enumeration<SparseFeatureVector> e : Iterables2.enumerate(data.unlabeledData())) {
			double[] dist = params.logJoint(e.getElement());
			initialY[e.getIndex()] =
					RandomGenerators.nextIntUnnormalizedLogProbs(rnd, dist);
		}
		return initialY;
	}
	
	// FIXME : should be "max" (or mean)
	private ProbabilisticModel sampleModel(CollapsedParameters sample) {
		double[] pOfY = DirichletDistribution.sample(sample.getTopicCounts(), rnd);
		double[][] pOfXGivenY = DirichletDistribution.sample(sample.getTopicWordCounts(), rnd);
		MoMParameters params = MoMParameters.fromProbabilities(pOfY, pOfXGivenY, false, false);
		return MixtureOfMultinomialsModel.newWithParameters(params, false);
	}
	
}
